import os
from ge.walker import RandomWalker
class GeneratorWalks():
    def __init__(self,data_save_root, walk_length, num_walks, p=1.0, q=1.0, workers=1, use_rejection_sampling=0):
        self.walk_length = walk_length
        self.num_walks = num_walks
        self.p = p
        self.q = q
        self.workers = workers
        self.use_rejection_sampling = use_rejection_sampling

        random_walk_file_name = f"steps={walk_length}-walks={num_walks}-return={q}-inout={p}.txt"
        self.save_file_name = os.path.join(data_save_root, random_walk_file_name)



    def load_walks(self,graph):
        if os.path.exists(self.save_file_name):
            sentences = []
            with open(self.save_file_name, 'r', encoding='utf-8') as fread:
                for line in fread:
                    sentences.append(line.strip().split(","))
        else:
            walker = RandomWalker(
                graph, p=self.p, q=self.q, use_rejection_sampling=self.use_rejection_sampling)

            print("Preprocess transition probs...")
            walker.preprocess_transition_probs()
            print("Preprocess random walks")
            sentences = walker.simulate_walks(
                num_walks=self.num_walks, walk_length=self.walk_length, workers=self.workers, verbose=1)

            with open(self.save_file_name,'w',encoding='utf-8') as fwriter:
                for item in sentences:
                    fwriter.write(",".join(item) +"\n")
        return sentences